package com.secsbrain.development.feign;

import com.netflix.client.config.IClientConfig;
import com.netflix.loadbalancer.ILoadBalancer;
import com.netflix.loadbalancer.reactive.LoadBalancerCommand;
import feign.Client;
import feign.Request;
import org.springframework.cloud.netflix.ribbon.ServerIntrospector;
import org.springframework.cloud.openfeign.ribbon.FeignLoadBalancer;

import java.net.URI;

public class SbFeignLoadBalancer extends FeignLoadBalancer {

    public SbFeignLoadBalancer(ILoadBalancer lb, IClientConfig clientConfig, ServerIntrospector serverIntrospector) {
        super(lb, clientConfig, serverIntrospector);
    }

    protected void customizeLoadBalancerCommandBuilder(final SbFeignLoadBalancer.SbRibbonRequest request, final IClientConfig config,
                                                       final LoadBalancerCommand.Builder<FeignLoadBalancer.RibbonResponse> builder) {
        request.setLoadBalancerKey();
    }

    protected static class SbRibbonRequest extends FeignLoadBalancer.RibbonRequest {

        protected SbRibbonRequest(Client client, Request request, URI uri) {
            super(client, request, uri);
        }

        void setLoadBalancerKey(){
            super.setLoadBalancerKey(getRequest().headers().get("loadBalancerKey"));
        }

    }

}
